import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks


def init_asr_pipeline(device="cuda:0"):
    """
    初始化ASR语音识别pipeline
    
    Args:
        device (str): 运行设备，默认为cuda:0
    
    Returns:
        pipeline: 初始化好的ASR pipeline
    """
    # 设置NPU设备
    torch_npu.npu.set_device("npu:0")
    
    # 初始化ASR pipeline
    asr_pipeline = pipeline(
        task=Tasks.auto_speech_recognition,
        model='model/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch',
        device=device
    )
    
    return asr_pipeline


if __name__ == "__main__":
    # 初始化pipeline
    asr_pipeline = init_asr_pipeline()
    print("ASR Pipeline initialized successfully!")